# -*- coding: utf-8 -*-

# @File  : data_util.py
# @Author: Lomo
# @Site  : lomo.space
# @Date  : 2020-02-12
# @Desc  : 数据读写


import os
import datetime

import csv

from config.config import PROJECT_DIR, DATA_DIR


class CSV(object):

    def __init__(self):
        self.file_path = PROJECT_DIR + DATA_DIR
        if not os.path.exists(self.file_path):
            os.mkdir(self.file_path)

    def create_csv(self, file_name=None):
        if not file_name:
            file_name = datetime.datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
        os.system('touch {0}/{1}.csv'.format(self.file_path, file_name))
        return '{0}/{1}.csv'.format(self.file_path, file_name)

    def read_csv(self, file_name):
        """
        读取 csv 文件并返回 list
        :param file_name: 文件名(默认为 /data 文件夹下)
        :return: list
        """
        abs_path = self.file_path + '/' + file_name + '.csv'
        if not os.path.exists(abs_path):
            return None
        with open(abs_path, encoding='utf-8') as f:
            csv_r = csv.reader(f)
            rows = [row for row in csv_r]
        return rows

    def write_csv(self, file_name, rows, row_headers=None):
        """
        以追加模式向 csv 文件写数据
        :param file_name: file name
        :param rows: 数据源格式支持 [] 或 [[], [], []] 或 [(), (), ()]
        :param row_headers: csv headers
        :return: None
        """
        abs_path = self.file_path + '/' + file_name + '.csv'
        print(abs_path)
        if self.read_csv(abs_path):
            has_header = self.read_csv(abs_path)[0] == row_headers
        else:
            has_header = False
        with open(abs_path, mode='a', encoding='utf-8') as f:
            f_csv = csv.writer(f)
            if row_headers is not None and has_header is False:
                f_csv.writerow(row_headers)
            if isinstance(rows, list):
                for obj in rows:
                    if isinstance(obj, (list, tuple)):
                        f_csv.writerow(list(obj))
                    else:
                        f_csv.writerow(rows)
                        break

    def del_file(self, file_name):
        abs_path = self.file_path + '/' + file_name + '.csv'
        if not os.path.exists(abs_path):
            return False
        else:
            try:
                os.system('rm {0}/{1}.csv'.format(self.file_path, file_name))
            except Exception as e:
                print(e)
            return True
